from typing import List, Optional, Protocol

from pydantic import BaseModel, Field
from volcenginesdkarkruntime.types.chat import ChatCompletionMessageParam

from arkitect.core.component.llm.model import (
    ArkChatParameters,
    ArkContextParameters,
    ChatCompletionTool,
)


class ToolType(Protocol):
    def tool_schema(self) -> ChatCompletionTool:
        ...


class State(BaseModel):
    model: str
    context_id: Optional[str] = Field(default=None)
    messages: List[ChatCompletionMessageParam] = Field(default_factory=list)
    parameters: Optional[ArkChatParameters] = Field(default=None)
    context_parameters: Optional[ArkContextParameters] = Field(default=None)
